# basic package
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np

# resize
from scipy import misc

# pytorch
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# set device
device = torch.device('cpu')

# hyper parameters
num_epochs = 5
num_classes = 10
batch_size = 32
learning_rate = 0.001

# cifar10 classes
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# data enhancement
transform = transforms.Compose([
    # padding to 36x36
    # transforms.Pad(4),
    # Random Horizontal Flip
    transforms.RandomHorizontalFlip(),
    # cut to 32x32
    transforms.RandomCrop(32),
    # to tensor
    transforms.ToTensor(),
    # normalization
    # transforms.Normalize(mean=(0.5, 0.5, 0.5),
    #                      std=(0.5, 0.5, 0.5))
    ])

# cifar10 path
cifar10Path = './cifar'

# train data
train_dataset = torchvision.datasets.CIFAR10(root=cifar10Path,
                                             train=True,
                                             transform=transform,
                                             download=True)

# test data
test_dataset = torchvision.datasets.CIFAR10(root=cifar10Path,
                                            train=False,
                                            transform=transform)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

data_iter = iter(test_loader)

images, labels = next(data_iter)

idx = 31
image = images[idx].numpy()
image = np.transpose(image, (1, 2, 0))
plt.imshow(image)
plt.show()
print(classes[labels[idx].numpy()])
print("1")
